'''
Created on Jan 16 2012

@author: oabalbin

This script takes
input: a list of snps positions by sample
output a file like CHR  POS    SP1   SP2   SP3
                    1   345    A/T   A/T   A/C
'''
import os
import sys
import subprocess
import numpy as np
from optparse import OptionParser
from collections import deque, defaultdict

                
def read_positions(ifile):
    '''
    '''
    ifile = open(ifile)
    for l in ifile:
        if l.startswith('#'):
            continue
        fields=l.strip('\n').split(',')
        fields.pop()
    
    ifile.close()
    return fields

def convert_thsnvs(thsnvs):
    '''
    chrM@2354|C>T|1/1,
    '''
    al1 = set()
    al2 = set()
    for snv in thsnvs:
        s = snv.split('|')
        p,a,b = s[0],s[1].split('>')[0],s[1].split('>')[1]
        al1.add( p+'|'+a+'>*' )
        al2.add( p+'|*>'+b )
    
    return al1,al2
    

def get_thsnvs_pos(thsnvs):
    '''
    chrM@2354|C>T|1/1,
    '''
    al1 = set()
    for snv in thsnvs:
        s = snv.split('|')
        p = s[0]
        al1.add( p )
    return al1

def get_ibd_snvs(pos_file,thsnvs):
    '''
    chrM@2354|C>T|1/1,
    '''
    ifile = open(pos_file)
    for l in ifile:
        if l.startswith('#'):
            continue
        fields=l.strip('\n').split(',')

    ibd_snvs=set()
    pos = get_thsnvs_pos(thsnvs)
    print "Getting the postions", len(pos)
    for i in fields:
        p0=i.split('|')[0]
        if p0 in pos:
            ibd_snvs.add( i )
    return ibd_snvs
    
    
def isec_arrays(filea,fileb):
    '''
    '''
    m1=set(read_positions(filea))
    m2=set(read_positions(fileb))
    #print sys.getsizeof(m1), sys.getsizeof(m2)
    isec=m1.intersection(m2)        
    return isec

def isec_snvs_annot_union(snvs_union, snvs_annot_dbs):
    '''
    This functions obtains all query snvs that are 
    common to a particular database, ex: dbSNP
    '''
    pm = defaultdict()
    for db, snv_db in snvs_annot_dbs.iteritems():
        pm[db] = set(snvs_union).intersection(set(snv_db))
           
    return pm

def union_snvs_dict(union_snvs):
    '''
    It gives an index to each snp. 
    So it can be used later to find arrays
    '''
    pm = defaultdict()
    for i, snv in enumerate(union_snvs):
        pm[snv]=i
    return pm


def read_patient_files(file):
    ifile=open(file)
    pat=defaultdict()
    
    for l in ifile:
        f=l.strip('\n').split('\t')
        pat[f[0]]=f[2]
    ifile.close()
    return pat


def ibd_children(patient_files,datadir):
    '''
    '''
    children = ['broth1','broth2','case']
    First=True
    for c in children:
        print os.path.join(datadir,patient_files[c])
        thsnvs = read_positions(os.path.join(datadir,patient_files[c]))
        # return 2 sets
        al1, al2 = convert_thsnvs(thsnvs)
        if First:
            al1_snvs_isec,al2_snvs_isec= al1,al2
            First=False
        else:
            al1_snvs_isec.intersection_update( al1 )
            al2_snvs_isec.intersection_update( al2 )
    return al1_snvs_isec,al2_snvs_isec


if __name__ == '__main__':
    '''
    Write to an array the positions in a particualr vcf.
    
    '''
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--pos_files", action="append", type="str",
                            dest="pos_files",
                            help="others_vcf_files. Pass vcf files without .gz extention even if the file .gz is there")
    
    optionparser.add_option("-p", "--patient_files", type="str", dest="patient_files",
                            help="others_vcf_files. Pass vcf files without .gz extention even if the file .gz is there")


    optionparser.add_option("-o", "--out_file", type="str",dest="ofile",
                            help="others_vcf_files. Pass vcf files without .gz extention even if the file .gz is there")
    
    optionparser.add_option("-b", "--base", dest="base", action="store_true", default=False,
                            help="determine if the ouput file is position or pos.base")
    optionparser.add_option("-g", "--geno", dest="geno", action="store_true", default=False,
                            help="determine if the ouput file is position or pos.base.geno")

    
    
    (options, args) = optionparser.parse_args()
    
    
    '''
    father    51    FMPN51.gatk_snps.b.filtered.model.pos.base.geno
    mother    65    FMPN65.gatk_snps.b.filtered.model.pos.base.geno
    broth1    52    FMPN52.gatk_snps.b.filtered.model.pos.base.geno
    case    63    FMPN63.gatk_snps.b.filtered.model.pos.base.geno
    broth2    53    FMPN53.gatk_snps.b.filtered.model.pos.base.geno
        
    '''
    
    patients_files = read_patient_files(options.patient_files)
    # Read pos.base.geno files
    First=True
    datadir='/exds/users/oabalbin/projects/fmpn/data/gatk_calls'
    # Get common positions for the children
    # Potential positions haploidentical parental
    al1_children, al2_children = ibd_children(patients_files,datadir)
    # Read the positions file and convert them for father
    thsnvs = read_positions(os.path.join(datadir,patients_files['father']))
    # return 2 sets
    al1, al2 = convert_thsnvs(thsnvs)
    al1_children.intersection_update( al1 )
    al2_children.intersection_update( al2 )
    
    # read the positions file and convert them for mother
    thsnvs = read_positions(os.path.join(datadir,patients_files['mother']))
    # return 2 sets
    al1, al2 = convert_thsnvs(thsnvs)
    al1_children.difference_update( al1 )
    al2_children.difference_update( al2 )
    
    # get the snvs in parental haplotype regions
    al1 = get_ibd_snvs(patients_files['case'],al1_children)
    al2 = get_ibd_snvs(patients_files['case'],al2_children)
    
    isec_set = al1.union(al2)
    
    print "The number of haploidentical parental postions is ",len(isec_set)
    # For the snvs shared across all individuals, isec set
    nrows=len(isec_set)
    ncols=1
    #mytype=[('CHRM',int),('POS',int),('IBD',int),('MUT','S3'),('GENO','S3'),('SAMPLE',int)]
    #mytype_ibd=[('CHRM',int),('POS',int),('IBD',int)]
    mytype=[('CHRM',int),('START',int),('END',int)]#,('INFO','S7')]
    mytype_ibd=[('CHRM',int),('START',int),('END',int)]#,('INFO','S7')]

    #location_arr=np.empty(( nrows, ncols),dtype=mytype)
    all_snvs=[]
    all_ibd=[]
    ibd=str(1)
    i=0
    for snv in isec_set:
        loc,base,geno = snv.split('|')[0],snv.split('|')[1].replace('>','/'),snv.split('|')[2]
        chr,pos=loc.split('@')[0],loc.split('@')[1]
        chrN=chr.replace('chr','')
        if chrN=='X':
            chrN=23
        elif chrN=='Y':
            chrN=24
        elif chrN=='M':
            chrN=25
        
        start, end =int(pos)-1,int(pos)+1
        info = base+'|'+geno
        all_snvs.append((chrN,start,end))#,info)) # This indixing is important, 0 means is shared for all samples
        all_ibd.append((chrN,start,end))#,info))

        #print ol, snv
        #print location_arr[i,:]
        i+=1
    

    # Sort the snv table
    location_arr=np.array( all_snvs,dtype=mytype )
    location_arr_sorted = np.sort(location_arr, order=['CHRM', 'START'])
    
    ########### Writing everything out
            
    # This takes time. Look for a better way to do it.    
    ibd_arr=np.array( all_ibd,dtype=mytype_ibd )
    ibd_arr_sorted=np.unique(ibd_arr)
    ibd_arr_sorted=np.sort(ibd_arr_sorted, order=['CHRM', 'START'])
    
    of=open(options.ofile,'w')   
    of2=open(options.ofile+'_all','w')
    
    # Write all the snps.
    chrs = np.array(range(1,26))
    i=1
    
    # Write all the snps.
    for i in range(len(location_arr_sorted)):
        
        ol=list(location_arr_sorted[i])
        if ol[0]==23:
            ol[0]='X'
        elif ol[0]==24:
            ol[0]='Y'
        elif ol[0]==25: # CHRM M
            continue
        #print ",".join(map(str,ol)).replace(',','\t')
        of2.write('chr'+",".join(map(str,ol)).replace(',','\t')+'\n')
    of2.close()           

    
    '''
    for i in chrs:
        sub_array=location_arr_sorted[location_arr_sorted['CHRM']==i]
        of2=open(options.ofile+'_chr'+str(i)+'_all.txt','w')
        print sub_array
        for j in range(len(sub_array)):
            ol=list(sub_array[j])
            if ol[0]==23:
                ol[0]='X'
            elif ol[0]==24:
                ol[0]='Y'
            elif ol[0]==25: # CHRM M
                continue
            #print ",".join(map(str,ol)).replace(',','\t')
            of2.write(",".join(map(str,ol)).replace(',','\t')+'\n')
        of2.close()
    '''
    '''
    # Write only the unique snps when ibd is considered
    for i in range(len(ibd_arr_sorted)):
        
        ol=list(ibd_arr_sorted[i])
        if ol[0]==23:
            ol[0]='X'
        elif ol[0]==24:
            ol[0]='Y'
        elif ol[0]==25: # CHRM M
            continue
        #print ",".join(map(str,ol)).replace(',','\t')
        of.write(",".join(map(str,ol)).replace(',','\t')+'\n')

    of.close()           
   '''